import shutil
import numpy as np
from tqdm import tqdm
from DDPM.Config import *
from scipy.spatial.distance import euclidean
from DDPM.Evaluation.Decoding import decode_all

def load_data_psfd(control_params, data_path: str, data_file_num:int, data_index_start:int=0):  # 载入PSFD数据
    from OPT.Objective import performance_matching_degree
    result_path = control_params["result_path"] 
    designed_npz_file = np.load(result_path + "/Optimization.npz")
    print(result_path)
    gv_designed = designed_npz_file["gv"][-1]
    lv_designed = designed_npz_file["lv"][-1]

    results = []  # 用于存储 (i, distance) 的列表

    for i in tqdm(range(data_index_start, data_index_start + data_file_num), desc='loading data', colour='green', dynamic_ncols=True):
        npz_path = f'{data_path}/{i}'
        npz_file = np.load(npz_path + ".npz")
        gv = npz_file["gv"]
        lv = npz_file["lv"]

        # 拼接 gv 和 lv 为 fv
        fv = np.concatenate((gv, lv), axis=0)
        fv_designed = np.concatenate((gv_designed, lv_designed), axis=0)

        # 计算 fv 和 fv_designed 之间的欧几里得距离
        distance = euclidean(fv, fv_designed)
        results.append((i, distance))

    # 按距离从小到大排序，并取前5个
    results.sort(key=lambda x: x[1])
    top_5_indices = [x[0] for x in results[:5]]
    print(results[:5])

    gvs=[]
    lvs=[]
    for idx in top_5_indices:
        npz_path = f'{data_path}/{idx}'
        npz_file = np.load(npz_path + ".npz")
        gv = npz_file["gv"]
        lv = npz_file["lv"]
        gvs.append(gv)
        lvs.append(lv)

    return lvs,gvs

def start_find_neighbor_in_psfd(control_params):
    result_path=control_params["result_path"] 
    # 清空目录
    if os.path.exists(f"{result_path}/Neighbor"):
        shutil.rmtree(f"{result_path}/Neighbor")
    os.makedirs(f"{result_path}/Neighbor", exist_ok=True)

    lvs,gvs=load_data_psfd(control_params, trainingset_path, data_file_num)

    # 保存PSFD中查询到的样本到新的 npz 文件
    save_path = f"{result_path}/Neighbor/Samples_neighbor.npz"
    np.savez_compressed(save_path,
                        lv=np.array(lvs),
                        gv=np.array(gvs))
    
    decode_all(f"{result_path}/Neighbor","Samples_neighbor")
    from DEAD.AutoDecoder.Evaluation.BurningSurface import burning_surface_all
    burning_surface_all(f"{result_path}/Neighbor")


if __name__ == "__main__":
    from OPT.Config import *

    define =make_define("dicT3",1,0.001,"LGA","Samples_clustered")
    control_params, dead_dic = get_control_params(define, False)

    start_find_neighbor_in_psfd(control_params)